"""
ConnLLM API密钥管理
负责管理LLM API密钥，支持从多个位置加载
"""
import os
import json
import logging
from pathlib import Path
from typing import Dict, Optional, Set, Union

from .constants import LOCAL_PROVIDERS
from .exceptions import ConfigurationError

# 设置日志
logger = logging.getLogger("connllm")

# 支持的提供商
SUPPORTED_PROVIDERS = LOCAL_PROVIDERS.union({
    "anthropic",    # Anthropic 提供商 (Claude模型)
    "openai",       # OpenAI 提供商 (GPT模型)
    "openrouter",   # OpenRouter 提供商 (聚合多个模型)
    "moonshot",     # Moonshot 提供商 (中文模型)
    "ppinfra",      # 鹏程立方 提供商
    "volces",       # Volces 提供商
    # 添加更多支持的提供商...
})


class KeyManager:
	"""API密钥管理器，支持从多个位置加载密钥"""
	
	def __init__(self, keys: Optional[Dict[str, str]] = None):
		"""
		初始化密钥管理器
		
		Args:
			keys: 密钥字典，提供商名称到API密钥的映射
		"""
		self.keys = {}
		
		# 尝试从多个位置加载密钥
		self._load_keys_from_files()
		
		# 使用传入的密钥覆盖
		if keys:
			self.update_keys(keys)
			
	def _load_keys_from_files(self) -> None:
		"""从多个位置加载密钥文件"""
		# 密钥文件搜索路径
		search_paths = self._get_key_search_paths()
		
		for path in search_paths:
			key_file = path / "connllm_keys.json"
			if key_file.exists():
				try:
					with open(key_file, "r", encoding="utf-8") as f:
						loaded_keys = json.load(f)
						
						# 验证格式
						if not isinstance(loaded_keys, dict):
							logger.warning(f"密钥文件格式错误: {key_file}")
							continue
							
						if "llm_keys" in loaded_keys and isinstance(loaded_keys["llm_keys"], dict):
							self.keys.update(loaded_keys["llm_keys"])
						else:
							self.keys.update(loaded_keys)
							
					logger.debug(f"从 {key_file} 加载了密钥")
				except (json.JSONDecodeError, IOError) as e:
					logger.warning(f"无法加载密钥文件 {key_file}: {str(e)}")
					
		# 从旧的 llm_keys.json 加载 (兼容性)
		for path in search_paths:
			key_file = path / "llm_keys.json"
			if key_file.exists():
				try:
					with open(key_file, "r", encoding="utf-8") as f:
						loaded_keys = json.load(f)
						
						# 验证格式
						if not isinstance(loaded_keys, dict):
							logger.warning(f"旧格式密钥文件格式错误: {key_file}")
							continue
							
						if "llm_keys" in loaded_keys and isinstance(loaded_keys["llm_keys"], dict):
							self.keys.update(loaded_keys["llm_keys"])
							
					logger.debug(f"从旧格式文件 {key_file} 加载了密钥")
				except (json.JSONDecodeError, IOError) as e:
					logger.warning(f"无法加载旧格式密钥文件 {key_file}: {str(e)}")
		
	def _get_key_search_paths(self) -> Set[Path]:
		"""
		获取密钥文件搜索路径
		
		Returns:
			Path集合
		"""
		paths = set()
		
		# 当前工作目录
		paths.add(Path.cwd())
		
		# 用户配置目录
		if os.name == "nt":  # Windows
			app_data = os.environ.get("APPDATA")
			if app_data:
				paths.add(Path(app_data) / "ConnLLM")
		else:  # Linux/MacOS
			home = os.environ.get("HOME")
			if home:
				paths.add(Path(home) / ".config" / "connllm")
				
		# 项目目录 (如果作为包安装)
		try:
			import connllm
			package_dir = Path(connllm.__file__).parent
			paths.add(package_dir)
		except ImportError:
			pass
			
		return paths
		
	def get_key(self, provider: str) -> Optional[str]:
		"""
		获取指定提供商的API密钥
		
		Args:
			provider: 提供商名称
			
		Returns:
			API密钥，如果不存在则返回None
			
		Raises:
			ConfigurationError: 如果提供商不受支持
		"""
		# 检查提供商是否受支持
		if provider not in SUPPORTED_PROVIDERS:
			logger.warning(f"未知的提供商类型: {provider}")
		
		# 本地提供商不需要API密钥
		if provider in LOCAL_PROVIDERS:
			return None
			
		# 1. 首先尝试从环境变量获取
		env_vars = [
			f"{provider.upper()}_API_KEY",  # 如: ANTHROPIC_API_KEY
			f"CONNLLM_{provider.upper()}_API_KEY"  # 如: CONNLLM_ANTHROPIC_API_KEY
		]
		
		for env_var in env_vars:
			if api_key := os.environ.get(env_var):
				return api_key
				
		# 2. 然后从内存中的密钥字典获取
		return self.keys.get(provider)
		
	def set_key(self, provider: str, api_key: Optional[str], save: bool = True) -> None:
		"""
		设置指定提供商的API密钥
		
		Args:
			provider: 提供商名称
			api_key: API密钥，如果为None则清除密钥
			save: 是否保存到文件
			
		Raises:
			ConfigurationError: 如果提供商不受支持
		"""
		# 检查提供商是否受支持
		if provider not in SUPPORTED_PROVIDERS:
			raise ConfigurationError(f"未知的提供商类型: {provider}")
			
		# 更新内存中的密钥
		if api_key is None and provider in self.keys:
			del self.keys[provider]
		elif api_key is not None:
			self.keys[provider] = api_key
			
		# 保存到文件
		if save:
			self._save_keys()
	
	def update_keys(self, keys: Dict[str, str], save: bool = True) -> None:
		"""
		批量更新密钥
		
		Args:
			keys: 密钥字典，提供商名称到API密钥的映射
			save: 是否保存到文件
		"""
		for provider, api_key in keys.items():
			if provider not in SUPPORTED_PROVIDERS:
				logger.warning(f"跳过未知提供商的密钥: {provider}")
				continue
				
			if api_key is None and provider in self.keys:
				del self.keys[provider]
			else:
				self.keys[provider] = api_key
				
		# 保存到文件
		if save:
			self._save_keys()
			
	def _save_keys(self) -> bool:
		"""
		保存密钥到用户配置目录
		
		Returns:
			是否成功保存
		"""
		# 用户配置目录
		if os.name == "nt":  # Windows
			app_data = os.environ.get("APPDATA")
			if app_data:
				user_config_dir = Path(app_data) / "ConnLLM"
			else:
				return False
		else:  # Linux/MacOS
			home = os.environ.get("HOME")
			if home:
				user_config_dir = Path(home) / ".config" / "connllm"
			else:
				return False
				
		# 创建用户配置目录
		try:
			user_config_dir.mkdir(parents=True, exist_ok=True)
		except Exception as e:
			logger.warning(f"无法创建用户配置目录: {str(e)}")
			return False
			
		# 密钥文件路径
		key_file = user_config_dir / "connllm_keys.json"
		
		try:
			# 使用新格式保存密钥
			with open(key_file, "w", encoding="utf-8") as f:
				json.dump(self.keys, f, indent=4, ensure_ascii=False)
				
			logger.debug(f"密钥已保存到 {key_file}")
			return True
		except Exception as e:
			logger.warning(f"无法保存密钥到文件 {key_file}: {str(e)}")
			return False
			
	def remove_key(self, provider: str, save: bool = True) -> bool:
		"""
		移除指定提供商的API密钥
		
		Args:
			provider: 提供商名称
			save: 是否保存到文件
			
		Returns:
			是否成功移除
		"""
		if provider in self.keys:
			del self.keys[provider]
			
			# 保存到文件
			if save:
				self._save_keys()
			return True
		return False
		
	def list_keys(self) -> Dict[str, Optional[str]]:
		"""
		获取所有可用的API密钥
		
		Returns:
			密钥字典，提供商名称到API密钥的映射
		"""
		result = {}
		
		# 本地提供商不需要API密钥
		for provider in SUPPORTED_PROVIDERS:
			if provider in LOCAL_PROVIDERS:
				result[provider] = None
			else:
				result[provider] = self.get_key(provider)
				
		return result
		
	def export_keys(self, path: Union[str, Path]) -> bool:
		"""
		导出所有密钥到文件
		
		Args:
			path: 导出文件路径
			
		Returns:
			是否成功导出
		"""
		try:
			with open(path, "w", encoding="utf-8") as f:
				json.dump(self.keys, f, indent=4, ensure_ascii=False)
				
			logger.debug(f"所有密钥已导出到 {path}")
			return True
		except Exception as e:
			logger.warning(f"无法导出密钥到文件 {path}: {str(e)}")
			return False
			
	def import_keys(self, path: Union[str, Path]) -> bool:
		"""
		从文件导入密钥
		
		Args:
			path: 导入文件路径
			
		Returns:
			是否成功导入
		"""
		try:
			with open(path, "r", encoding="utf-8") as f:
				loaded_keys = json.load(f)
				
			if not isinstance(loaded_keys, dict):
				logger.warning(f"密钥文件格式错误: {path}")
				return False
				
			# 处理旧格式
			if "llm_keys" in loaded_keys and isinstance(loaded_keys["llm_keys"], dict):
				loaded_keys = loaded_keys["llm_keys"]
				
			# 更新密钥
			for provider, api_key in loaded_keys.items():
				if provider not in SUPPORTED_PROVIDERS:
					logger.warning(f"跳过未知提供商的密钥: {provider}")
					continue
					
				self.keys[provider] = api_key
				
			# 保存到文件
			self._save_keys()
			
			logger.debug(f"从 {path} 导入了密钥")
			return True
		except Exception as e:
			logger.warning(f"无法从文件 {path} 导入密钥: {str(e)}")
			return False
			
	def validate_key(self, provider: str, api_key: Optional[str] = None) -> bool:
		"""
		验证API密钥是否有效
		
		Args:
			provider: 提供商名称
			api_key: 要验证的API密钥，如果为None则验证当前密钥
			
		Returns:
			密钥是否有效
		"""
		# 本地提供商不需要API密钥
		if provider in LOCAL_PROVIDERS:
			return True
			
		# 使用指定的密钥或当前密钥
		if api_key is None:
			api_key = self.get_key(provider)
			
		if not api_key:
			logger.warning(f"未提供{provider}的API密钥")
			return False
			
		# 验证密钥格式
		if not self._validate_key_format(provider, api_key):
			logger.warning(f"{provider}的API密钥格式无效")
			return False
			
		# 尝试使用API密钥进行简单调用
		# 注意：这会消耗API使用量，未实现
		
		return True
		
	def _validate_key_format(self, provider: str, api_key: str) -> bool:
		"""
		验证API密钥格式
		
		Args:
			provider: 提供商名称
			api_key: API密钥
			
		Returns:
			格式是否有效
		"""
		# Anthropic API密钥格式: sk-ant-api03-...
		if provider == "anthropic" and api_key.startswith("sk-ant-"):
			return True
			
		# OpenAI API密钥格式: sk-...
		if provider == "openai" and api_key.startswith("sk-"):
			return True
			
		# OpenRouter API密钥格式: sk-or-...
		if provider == "openrouter" and api_key.startswith("sk-or-"):
			return True
			
		# Moonshot API密钥格式: sk-ms-...
		if provider == "moonshot" and api_key.startswith("sk-ms-"):
			return True
			
		# 如果没有特定格式验证，返回True
		return True
